#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <openssl/rsa.h>
#include <openssl/pem.h>
#include <openssl/err.h>

#define RSA_KEY_BITS 2048

void generateKeys(){
    EVP_PKEY* pkey = EVP_RSA_gen(1024);
    if(pkey == NULL){
        fprintf(stderr,"error: rsa gen\n");
        ERR_print_errors_fp(stderr);
        return;
    }
    FILE* fp = fopen("public.pem","wt");
    if(fp!= NULL){
        PEM_write_PUBKEY(fp,pkey);
        fclose(fp);
    }else{
        perror("file error");
    }
    fp = fopen("private.pem","wt");
    if(fp != NULL){
        PEM_write_PrivateKey(fp,pkey,NULL,NULL,0,NULL,NULL);
        fclose(fp);
    }else{
        perror("file error");
    }
    EVP_PKEY_free(pkey);
}

char *encrypt_public(char *src, int len, int *length)
{
    FILE *fp = fopen("public.pem", "r");
    if (fp == NULL)
    {
        perror("file error");
        return NULL;
    }
    EVP_PKEY *pkey;
    pkey = PEM_read_PUBKEY(fp, NULL, NULL, NULL);
    fclose(fp);
    if (pkey == NULL)
    {
        fprintf(stderr, "error: read publics key\n");
        return NULL;
    }
    EVP_PKEY_CTX *ctx = EVP_PKEY_CTX_new(pkey, NULL);
    EVP_PKEY_encrypt_init(ctx);
    char *dst = (char *)malloc(2048);
    size_t outl = 2048;
    if (!EVP_PKEY_encrypt(ctx, dst, &outl, src, (size_t)len))
    {
        fprintf(stderr, "error: encrypt\n");
        EVP_PKEY_free(pkey);
        free(dst);
        return NULL;
    }
    int len2 = outl;
    EVP_PKEY_free(pkey);
    EVP_PKEY_CTX_free(ctx);
    BIO_dump_fp(stdout, dst, len2);
    printf("len: %d, len2: %d\n", len, len2);
    if (length != NULL)
    {
        *length = len2;
    }
    return dst;
}

char *decrypt_private(char *src, int len, int plaintext_len)
{
    FILE *fp = fopen("private.pem", "r");
    if (fp == NULL)
    {
        perror("file error");
        return NULL;
    }
    EVP_PKEY *pkey = PEM_read_PrivateKey(fp, NULL, NULL, NULL);
    fclose(fp);
    if (pkey == NULL)
    {
        fprintf(stderr, "error: read private key\n");
        return NULL;
    }
    EVP_PKEY_CTX *ctx = EVP_PKEY_CTX_new(pkey, NULL);
    EVP_PKEY_decrypt_init(ctx);
    char *dst = (char *)malloc(plaintext_len);
    size_t outl;
    size_t inl = len;
    if (!EVP_PKEY_decrypt(ctx, dst, &outl, src, inl))
    {
        fprintf(stderr, "error: decrypt\n");
        free(dst);
        dst = NULL;
    }
    else
    {
        BIO_dump_fp(stdout, dst, (int)outl);
        printf("len: %d, outl: %ld\n", len, outl);
    }
    EVP_PKEY_free(pkey);
    EVP_PKEY_CTX_free(ctx);
    return dst;
}



int main() {
    // Generate RSA keys
    generateKeys();

    // Test encryption and decryption
    char *plaintext = "hello world nihao!ww";
    int plaintext_len = strlen(plaintext);
    printf("Plaintext: %s\n", plaintext);

    // Encrypt with public key
    int ciphertext_len;
    char *ciphertext = encrypt_public(plaintext, plaintext_len, &ciphertext_len);
    if (ciphertext == NULL) {
        fprintf(stderr, "Error encrypting.\n");
        return 1;
    }
    printf("Ciphertext:\n");
    for (int i = 0; i < ciphertext_len; i++) {
        printf("%02X ", (unsigned char)ciphertext[i]);
    }
    printf("\n");

    // Decrypt with private key
   char *decryptedtext = decrypt_private(ciphertext, ciphertext_len, plaintext_len);
    if (decryptedtext == NULL) {
        fprintf(stderr, "Error decrypting.\n");
        free(ciphertext);
        return 1;
    }
    printf("Decrypted text: %s\n", decryptedtext);

    // Clean up
    free(ciphertext);
    free(decryptedtext);

    return 0;
}